module basic

    use, intrinsic :: iso_c_binding
    implicit none
    include "fftw3.fi"
    real(8), parameter :: pi2 = 6.283185307179586D0

    interface lstsq
        module procedure leastsqs, leastsqm
    end interface ! lstsq

contains

function nextpow2(n) bind(c)

    integer, intent(in) :: n
    integer :: nextpow2

    nextpow2 = 1

    do while ( nextpow2<n )
        nextpow2 = nextpow2*2
    end do
    return

end function nextpow2

function nextpow(n,base) bind(c)

    integer, intent(in) :: n, base
    integer :: nextpow

    nextpow = 1

    do while ( nextpow<n )
        nextpow = nextpow*base
    end do
    return

end function nextpow

function mean(x,n) bind(c)
!     平均值
    integer, intent(in) :: n
    real(8), intent(in) :: x(n)
    real(8) :: mean

    mean = sum(x)/dble(n)

end function mean

function rms(x,n) bind(c)
!     均方根值
    integer, intent(in) :: n
    real(8), intent(in) :: x(n)
    real(8) :: rms

    rms = sqrt(sum(x*x)/dble(n))
end function rms

subroutine findpeak(x,n,p,i) bind(c)
!     查找峰值及其索引
    integer, intent(in) :: n
    real(8), intent(in) :: x(n)
    real(8), intent(out) :: p
    integer, intent(out) :: i

    integer, external :: idamax

    i = idamax(n,x,1);
    p = x(i);

    return
end subroutine findpeak

function peak(x,n) bind(c)
!     峰值
    
    integer, intent(in) :: n
    real(8), intent(in) :: x(n)
    real(8) :: peak

    integer :: i

    call findpeak(x,n,peak,i)

    return

end function peak

subroutine norm(a,n) bind(c)
!     按峰值归一
    integer, intent(in) :: n
    real(8), intent(inout) :: a(n)

    integer :: i
    real(8) :: pk, ai

    pk = 0.d0

    do i = 1, n, 1
        ai = abs(a(i))
        if ( ai>pk ) pk = ai
    end do

    !$OMP PARALLEL DO
    do i = 1, n, 1
        a(i) = a(i)/pk
    end do
    !$OMP END PARALLEL DO

    return

end subroutine norm

function dpower(x,n,m) bind(c)
!     对x的n次幂求m阶导数，即diff(x^n,x,m)
    real(8), intent(in) :: x
    integer, intent(in) :: n, m
    real(8) :: dpower

    integer :: k, l

    if ( m>n ) then
        dpower = 0.d0
        return
    end if

    k = n
    l = m
    dpower = x**(n-m)

    do while ( l>0 )
        dpower = dpower*k
        l = l - 1
        k = k - 1
    end do

end function dpower

function trapz(y,n,dx,y0) bind(c)
!     梯形法求数值积分，离散点，间隔均匀(dx)
    implicit none
    real*8, intent(in) :: dx, y0
    integer, intent(in) :: n
    real*8, intent(in) :: y(n)

    real*8 :: trapz

    integer :: i
    real*8 :: r

    r = 0.d0

    !$OMP PARALLEL PRIVATE(i), SHARED(y,dx,n)
    !$OMP DO REDUCTION(+:r)

    do i = 2, n, 1
        r = r + dx*(y(i)+y(i-1))
    end do
    !$OMP END DO
    !$OMP END PARALLEL

    trapz = y0+r*0.5
    return

end function trapz

subroutine cumtrapz(y,z,n,dx,z0) bind(c)
!     梯形法求累积数值积分，离散点，间隔均匀(dx)
    real*8, intent(in) :: dx, z0
    integer, intent(in) :: n
    real*8, intent(in) :: y(n)
    real*8, intent(out) :: z(n)

    integer :: i
    real*8 :: r

    z(1) = z0
    r = 0.D0
    do i = 2, n, 1
        r = r + dx*(y(i)+y(i-1))*0.5
        z(i) = z0 + r
    end do
    return

end subroutine cumtrapz

subroutine ariasIntensity(a,Ia,n) bind(c)
!     归一化累积 Arias 强度
    integer, intent(in) :: n
    real*8, intent(in) :: a(n)
    real*8, intent(out) :: Ia(n)

    call cumtrapz(a*a,Ia,n,1.d0,0.d0)
    Ia = Ia/Ia(n)
    return

end subroutine ariasIntensity

subroutine fftfreqs(Nfft,fs,freqs) bind(c)
!     FFT 变换对应的频率坐标

    integer, intent(in) :: Nfft
    real(8), intent(in) :: fs
    real(8), intent(out) :: freqs(Nfft)

    real(8) :: fn, df
    integer :: i, l

    fn = 0.5D0*fs
    df = fs/dble(Nfft)
    freqs = 0.0D0

    if ( mod(Nfft, 2) == 0 ) then
        l = Nfft/2-1
        freqs(2:Nfft/2) = [ ( (dble(i)*df), i=1,l ) ]
        freqs(Nfft/2+1:Nfft) = [ ( -fn+(dble(i-1)*df), i=1,l+1 ) ]
    else
        l = (Nfft+1)/2-1
        freqs(2:(Nfft+1)/2) = [ ( (dble(i)*df), i=1,l ) ]
        freqs((Nfft+1)/2+1:Nfft) = [ ( -(dble(l-i+1)*df), i=1,l ) ]
    end if

    return

end subroutine fftfreqs

subroutine leastsqs(a,b,m,n) bind(c)
!     Ax = b 的最小二乘解，A 为 m × n 阶矩阵，
!     b 为 n 阶向量，结果保存在 b 中
    integer, intent(in) :: m, n
    real*8, intent(in) :: a(:,:)
    real*8, intent(inout) :: b(:)

    integer :: i, nrhs, lda, ldb, lwork
    integer :: r, info
    real*8 :: twork(1)
    integer :: tiwork(1)
    real*8, allocatable :: s(:)
    real*8, allocatable :: work(:)
    integer,allocatable :: iwork(:)

    ! m = size(a, dim=1)
    ! n = size(a, dim=2)
    nrhs = 1

    lda = max(1,m)
    ldb = max(1,m,n)

    allocate(s(min(m,n)))

    lwork = -1

    call dgelsd(m, n, nrhs, a, lda, b, ldb, s, -1.d0, &
        r, twork, lwork, tiwork, info)

    lwork = int(twork(1))
    allocate(work(lwork))
    allocate(iwork(tiwork(1)))

    call dgelsd(m, n, nrhs, a, lda, b, ldb, s, -1.d0, &
        r, work, lwork, iwork, info)

    deallocate(s)
    deallocate(work)
    deallocate(iwork)

    return

end subroutine leastsqs

subroutine leastsqm(a,b,m,n,nrhs) bind(c)
!     Ax = b 的最小二乘解，A 为 m × n 阶矩阵
!     b 为 n × nrhs 阶矩阵，结果保存在 b 中
    real*8, intent(in) :: a(:,:)
    real*8, intent(inout) :: b(:,:)
    integer, intent(in) :: m, n, nrhs

    integer :: i, lda, ldb, lwork
    integer :: rank, info
    real*8 :: twork(1)
    integer :: tiwork(1)
    real*8, allocatable :: s(:)
    real*8, allocatable :: work(:)
    integer,allocatable :: iwork(:)

    ! m = size(a, dim=1)
    ! n = size(a, dim=2)
    ! nrhs = size(b, dim=2)

    lda = max(1,m)
    ldb = max(1,m,n)

    allocate(s(min(m,n)))

    lwork = -1

    call dgelsd(m, n, nrhs, a, lda, b, ldb, s, -1.d0, &
        rank, twork, lwork, tiwork, info)

    lwork = int(twork(1))
    allocate(work(lwork))
    allocate(iwork(tiwork(1)))

    call dgelsd(m, n, nrhs, a, lda, b, ldb, s, -1.d0, &
        rank, work, lwork, iwork, info)

    deallocate(s)
    deallocate(work)
    deallocate(iwork)

    return

end subroutine leastsqm

subroutine error(y,y0,n,aerror,merror) bind(c)
!     计算y和y0之间的平均误差和最大误差（不考虑首尾的数值）
    integer, intent(in) :: n
    real(8), intent(in) :: y(:), y0(:)
    real(8), intent(out) :: aerror, merror

    real(8), allocatable :: e(:)

    allocate(e(n-2))

    e = (y(2:n-1) - y0(2:n-1))/y0(2:n-1)

    aerror = sqrt(sum(e*e)/dble(n-2))
    merror = maxval(abs(e), dim=1)

    deallocate(e)

end subroutine error

subroutine errora(y,y0,n,aerror,merror) bind(c)
!     计算y和y0之间的平均误差和最大误差
    integer, intent(in) :: n
    real(8), intent(in) :: y(:), y0(:)
    real(8), intent(out) :: aerror, merror

    real(8), allocatable :: e(:)

    allocate(e(n))

    e = (y - y0)/y0

    aerror = sqrt(sum(e*e)/dble(n))
    merror = maxval(abs(e), dim=1)

    deallocate(e)

end subroutine errora

subroutine incrlininterp(x,y,n,xi,yi,ni) bind(c)
!     x 为升序排列时对 y 进行线性插值
    integer, intent(in) :: n, ni
    real(8), intent(in) :: x(:), y(:), xi(:)
    real(8), intent(out) :: yi(:)

    integer :: i, j
    real(8) :: xc, yc, xp, yp, slope

    xp = x(1)
    yp = y(1)
    j = 1
    do i = 2, n, 1
        xc = x(i)
        yc = y(i)
        slope = (yc - yp)/(xc - xp)
        do while ( xi(j)<=xc .and. j<=ni)
            yi(j) = yp + slope*(xi(j)-xp)
            j = j + 1
        end do
        xp = xc
        yp = yc
    end do

end subroutine incrlininterp

subroutine incrfindfirst(a,n,x,i) bind(c)
!     a为升序排列时寻找a中第一个超过x的数的索引
    integer, intent(in) :: n
    real(8), intent(in) :: a(:), x
    integer, intent(out) :: i

    i = 1
    do while ( a(i)<x .and. i<=n )
        i = i + 1
    end do

    return

end subroutine incrfindfirst

subroutine incrfindlast(a,n,x,i) bind(c)
!     a为升序排列时，从a的尾端开始寻找a中第一个小于x的数的索引
    integer, intent(in) :: n
    real(8), intent(in) :: a(:), x
    integer, intent(out) :: i

    i = n
    do while ( a(i)>x .and. i>=1 )
        i = i - 1
    end do

    return

end subroutine incrfindlast

subroutine decrlininterp(x,y,n,xi,yi,ni) bind(c)
!     x 为降序排列时对 y 进行线性插值
    integer, intent(in) :: n, ni
    real(8), intent(in) :: x(:), y(:), xi(:)
    real(8), intent(out) :: yi(:)

    integer :: i, j
    real(8) :: xc, yc, xp, yp, slope

    xp = x(n)
    yp = y(n)
    j = 1
    do i = 1, n-1, 1
        xc = x(n-i)
        yc = y(n-i)
        slope = (yc - yp)/(xc - xp)
        do while ( xi(j)>=xc .and. j<=ni)
            yi(j) = yp + slope*(xi(j)-xp)
            j = j + 1
        end do
        xp = xc
        yp = yc
    end do

end subroutine decrlininterp

subroutine decrfindfirst(a,n,x,i) bind(c)
!     a为降序排列时，寻找a中第一个小于x的数的索引
    integer, intent(in) :: n
    real(8), intent(in) :: a(:), x
    integer, intent(out) :: i

    i = 1
    do while ( a(i)>x .and. i<=n )
        i = i + 1
    end do

    return

end subroutine decrfindfirst

subroutine decrfindlast(a,n,x,i) bind(c)
!     a为降序排列时，从a的尾端开始寻找a中第一个超过x的数的索引
    integer, intent(in) :: n
    real(8), intent(in) :: a(:), x
    integer, intent(out) :: i

    i = n
    do while ( a(i)<x .and. i>=1 )
        i = i - 1
    end do

    return

end subroutine decrfindlast

subroutine fftpadding(in,out,n,nfft) bind(c)
!     数据序列后补零
    integer, intent(in) :: n, nfft
    real*8, intent(in) :: in(n)
    real*8, intent(out) :: out(nfft)

    out(1:n) = in(1:n)
    out(n+1:nfft) = 0.d0

    return
end subroutine fftpadding

subroutine fft(in,out,n) bind(c)
!     FFT 变换 实数->复数

    integer, intent(in) :: n
    real*8, intent(inout) :: in(n)
    complex*16, intent(inout) :: out(n)

    type(c_ptr) :: plan

    plan = fftw_plan_dft_r2c_1d(n, in, out, fftw_estimate)
    call fftw_execute_dft_r2c(plan, in, out)
    call fftw_destroy_plan(plan)

    return
end subroutine fft

subroutine ifft(in,out,n) bind(c)
!     FFT 逆变换 复数->实数

    integer, intent(in) :: n
    complex*16, intent(inout) :: in(n)
    real*8, intent(inout) :: out(n)

    type(c_ptr) :: plan

    plan = fftw_plan_dft_c2r_1d(n, in, out, fftw_estimate)
    call fftw_execute_dft_c2r(plan, in, out)
    call fftw_destroy_plan(plan)
    out = out/dble(n)

    return
end subroutine ifft

subroutine hann(w,n) bind(c)
!     对 w 施加 hann 窗

    integer(C_INT), intent(in) :: n
    real(C_DOUBLE), intent(inout) :: w(n)

    integer(C_INT) i

    do i = 1, n, 1
        w(i) = 2.d0*w(i)*(0.5d0 - 0.5d0 * cos(pi2 * dble(i - 1) / dble(n - 1)))
    end do

end subroutine hann

subroutine welch(a,n,m,olr,psd,win) bind(c)
!     Welch 法估计功率谱密度
    integer(C_INT), intent(in) :: n, m
    real(C_DOUBLE), intent(in) :: a(n)
    real(C_DOUBLE), intent(out) :: psd(m)
    real(C_DOUBLE), intent(in) :: olr
    integer(C_INT), intent(in) :: win

    integer(C_INT) :: nol, i, j, k, q, nsub
    real(C_DOUBLE) :: re, im
    real(C_DOUBLE), allocatable :: c(:)
    complex(C_DOUBLE_COMPLEX), allocatable :: cf(:)

    type(c_ptr) :: plan

    if ( mod(m,2) == 0 ) then
        q = m/2 + 1
    else
        q = (m+1)/2
    end if

    nol = int(m*olr)
    k = m
    nsub = 1
    do while ( k < n )
        k = k - nol + m
        nsub = nsub + 1
    end do

    allocate(c(m))
    allocate(cf(m))

    psd = 0.d0
    c = a(1:m)
    plan = fftw_plan_dft_r2c_1d(m, c, cf, fftw_estimate)
    if (win>0) then 
        call hann(c,m)
    end if
    call fftw_execute_dft_r2c(plan, c, cf)

    do j = 1, q, 1
        re = dble(cf(j))
        im = aimag(cf(j))
        psd(j) = (re*re + im*im)/dble(m*nsub)
    end do

    k = m
    do i = 2, nsub-1, 1
        k = k - nol + m
        c = a((k-m+1):k)
        if (win>0) then 
            call hann(c,m)
        end if
        call fftw_execute_dft_r2c(plan, c, cf)
        do j = 1, q, 1
            re = dble(cf(j))
            im = aimag(cf(j))
            psd(j) = psd(j) + (re*re + im*im)/dble(m*nsub)
        end do
    end do

    k = k - nol + m
    c(1:(n-k+m)) = a((k-m+1):n)
    c((n-k+m+1):m) = 0.d0
    if (win>0) then 
        call hann(c,m)
    end if
    call fftw_execute_dft_r2c(plan, c, cf)
    
    do j = 1, q, 1
        re = dble(cf(j))
        im = aimag(cf(j))
        psd(j) = psd(j) + (re*re + im*im)/dble(m*nsub)
    end do

    call fftw_destroy_plan(plan)

    deallocate(c)
    deallocate(cf)

end subroutine welch

subroutine fftresample(a,n,r,ar,nr) bind(c)

    integer(C_INT), intent(in) :: n, r, nr
    real(C_DOUBLE), intent(in) :: a(n)
    real(C_DOUBLE), intent(out) :: ar(nr)

    integer(C_INT) :: Nfft, Nfftr
    real(C_DOUBLE), allocatable :: ap(:), apr(:)
    complex(C_DOUBLE_COMPLEX), allocatable :: af(:), afr(:)

    if ( r == 1 ) return

    Nfft = nextpow(r,n)

    allocate(ap(Nfft))
    allocate(af(Nfft))

    call fftpadding(a,ap,n,Nfft)
    call fft(ap,af,Nfft)

    Nfftr = Nfft/r
    allocate(apr(Nfft))
    allocate(afr(Nfft))

    if ( mod(Nfftr,2) == 0 ) then
        afr(1:Nfftr/2) = af(1:Nfftr/2)
        afr(Nfftr/2+1) = 1.0*af(Nfftr/2+1)
    else
        afr(1:(Nfftr-1)/2) = af(1:(Nfftr-1)/2)
        afr((Nfftr-1)/2+1) = 1.0*af((Nfftr-1)/2+1)
    end if

    call ifft(afr,apr,Nfftr)

    ar = (1.d0/dble(r))*apr(1:nr)

    deallocate(ap)
    deallocate(af)
    deallocate(apr)
    deallocate(afr)

    return
    
end subroutine fftresample

end module basic